#include "linefit.h"

LineFit::LineFit() {}

LineFit::~LineFit() {}

Eigen::Vector2f LineFit::LeastSquareFitBeeline(const std::vector<float>& X,
                                               const std::vector<float>& Y) {
  if (X.size() != Y.size() || X.size() < 2 || Y.size() < 2) {
    exit(EXIT_FAILURE);
  }
  std::chrono::steady_clock::time_point t1 = std::chrono::steady_clock::now();

  int n = X.size();
  float sum_xy = 0.0f;
  float sum_x = 0.0f;
  float sum_y = 0.0f;
  float sum_xx = 0.0f;
  for (int i = 0; i < n; ++i) {
    sum_x += X[i];
    sum_y += Y[i];
    sum_xx += X[i] * X[i];
    sum_xy += X[i] * Y[i];
  }
  float denominator = n * sum_xx - sum_x * sum_x;
  if (fabs(denominator) <= 1e-3) denominator = 1e-3;
  // y = kx+b
  float K = (n * sum_xy - sum_x * sum_y) / denominator;
  float B = (sum_y * sum_xx - sum_x * sum_xy) / denominator;

  std::chrono::steady_clock::time_point t2 = std::chrono::steady_clock::now();
  std::chrono::microseconds time_used =
      std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1);
  printf("time used [%s] : %ld ms\n", __FUNCTION__, time_used.count());

  return Eigen::Vector2f(B, K);
}
Eigen::Vector3f LineFit::LeastSquareFitCurve(const std::vector<float>& X,
                                             const std::vector<float>& Y) {
  if (X.size() != Y.size() || X.size() < 2 || Y.size() < 2) {
    exit(EXIT_FAILURE);
  }
  std::chrono::steady_clock::time_point t1 = std::chrono::steady_clock::now();
  int n = X.size();
  float sum_xy = 0.0f;
  float sum_x = 0.0f;
  float sum_y = 0.0f;
  float sum_xx = 0.0f;

  float sum_x4 = 0.0f;
  float sum_x3 = 0.0f;
  float sum_x2y = 0.0f;
  for (int i = 0; i < n; ++i) {
    sum_x += X[i];
    sum_y += Y[i];
    sum_xx += X[i] * X[i];
    sum_xy += X[i] * Y[i];
    sum_x4 += std::pow(X[i], 4);
    sum_x3 += pow(X[i], 3);
    sum_x2y += X[i] * X[i] * Y[i];
  }
  float denominator = pow(sum_xx, 3) - 2 * sum_xx * sum_x3 * sum_x -
                      sum_xx * sum_x4 * n + sum_x3 * sum_x3 * n +
                      sum_x4 * sum_x * sum_x;
  if (fabs(denominator) <= 1e-3) denominator = 1e-3;
  float a = -sum_x2y * (n * sum_xx - sum_x * sum_x) -
            sum_xy * (sum_xx * sum_x - sum_x3 * n) +
            sum_y * (sum_xx * sum_xx - sum_x3 * sum_x);
  float b = -sum_x2y * (sum_x * sum_xx - sum_x3 * n) +
            sum_xy * (sum_xx * sum_xx - sum_x4 * n) -
            sum_y * (sum_xx * sum_x3 - sum_x4 * sum_x);
  float c = sum_x2y * (sum_xx * sum_xx - sum_x3 * sum_x) -
            sum_xy * (sum_xx * sum_x3 - sum_x4 * sum_x) -
            sum_y * (sum_xx * sum_x4 - sum_x3 * sum_x3);

  std::chrono::steady_clock::time_point t2 = std::chrono::steady_clock::now();
  std::chrono::microseconds time_used =
      std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1);
  printf("time used [%s] : %ld ms\n", __FUNCTION__, time_used.count());
  // y = ax²+bx+c
  c = c / denominator;
  b = b / denominator;
  a = a / denominator;
  return Eigen::Vector3f(c, b, a);
}
Eigen::VectorXf LineFit::SolveBySVD(const Eigen::MatrixXf& A,
                                    const Eigen::MatrixXf& b) {
  using namespace Eigen;
  return A.bdcSvd(ComputeThinU | ComputeThinV).solve(b);
}
Eigen::VectorXf LineFit::SolveByQR(const Eigen::MatrixXf& A,
                                   const Eigen::MatrixXf& b) {
  return A.colPivHouseholderQr().solve(b);
}
Eigen::VectorXf LineFit::SolveByLDLT(const Eigen::MatrixXf& A,
                                     const Eigen::MatrixXf& b) {
// If the matrix A is ill-conditioned, then this is not a good method,
// because the condition number of ATA is the square of the condition number of
// A. This means that you lose twice as many digits using normal equation than
// if you use the other methods.
  return (A.transpose() * A).ldlt().solve(A.transpose() * b);
}
/**
 * @description: Ax=b，构建矩阵A和b
 * @param {XY 为对应坐标系的点坐标}
 * @param {orders 代表几次方程，为2时求解二项式曲线y=t0+t1*x+t2*x^2}
 * @return {*}
 */
void LineFit::ConstructAb(const std::vector<float>& X,
                          const std::vector<float>& Y, uint8_t orders,
                          Eigen::MatrixXf& A, Eigen::MatrixXf& b) {
  // abnormal input verification
  if (X.size() < 2 || Y.size() < 2 || X.size() != Y.size() || orders < 1)
    exit(EXIT_FAILURE);

  // map sample data from STL vector to eigen vector
  std::vector<float> x_copy(X), y_copy(Y);
  Eigen::Map<Eigen::VectorXf> sampleX(x_copy.data(), X.size());
  Eigen::Map<Eigen::VectorXf> sampleY(y_copy.data(), Y.size());

  // Vandermonde matrix of X-axis coordinate vector of sample data
  Eigen::MatrixXf mtxVandermonde(X.size(), orders + 1);

  // Vandermonde column
  Eigen::VectorXf colVandermonde = sampleX;

  // construct Vandermonde matrix column by column
  for (size_t i = 0; i < orders + 1; ++i) {
    if (0 == i) {
      mtxVandermonde.col(0) = Eigen::VectorXf::Constant(X.size(), 1, 1);
      continue;
    }
    if (1 == i) {
      mtxVandermonde.col(1) = colVandermonde;
      continue;
    }
    colVandermonde = colVandermonde.array() * sampleX.array();
    mtxVandermonde.col(i) = colVandermonde;
  }

  // calculate coefficients vector of fitted polynomial
  A = mtxVandermonde.transpose() * mtxVandermonde;
  b = mtxVandermonde.transpose() * sampleY;
}
Eigen::VectorXf LineFit::LeastSquareNormalMethod(const std::vector<float>& X,
                                                 const std::vector<float>& Y,
                                                 uint8_t orders) {
  std::chrono::steady_clock::time_point t1 = std::chrono::steady_clock::now();

  ConstructAb(X, Y, orders, A_, b_);
  Eigen::VectorXf result = A_.inverse() * b_;

  std::chrono::steady_clock::time_point t2 = std::chrono::steady_clock::now();
  std::chrono::microseconds time_used =
      std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1);
  printf("time used [%s] : %ld ms\n", __FUNCTION__, time_used.count());

  return result;
}

Eigen::VectorXf LineFit::LeastSquareSVDMethod(const std::vector<float>& X,
                                              const std::vector<float>& Y,
                                              uint8_t orders) {
  std::chrono::steady_clock::time_point t1 = std::chrono::steady_clock::now();

  ConstructAb(X, Y, orders, A_, b_);

  std::chrono::steady_clock::time_point t2 = std::chrono::steady_clock::now();
  std::chrono::microseconds time_used =
      std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1);
  printf("time used [%s] : %ld ms\n", __FUNCTION__, time_used.count());

  return SolveBySVD(A_, b_);
}
Eigen::VectorXf LineFit::LeastSquareQRMethod(const std::vector<float>& X,
                                             const std::vector<float>& Y,
                                             uint8_t orders) {
  std::chrono::steady_clock::time_point t1 = std::chrono::steady_clock::now();

  ConstructAb(X, Y, orders, A_, b_);

  std::chrono::steady_clock::time_point t2 = std::chrono::steady_clock::now();
  std::chrono::microseconds time_used =
      std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1);
  printf("time used [%s] : %ld ms\n", __FUNCTION__, time_used.count());

  return SolveByQR(A_, b_);
}

Eigen::VectorXf LineFit::LeastSquareLDLTMethod(const std::vector<float>& X,
                                               const std::vector<float>& Y,
                                               uint8_t orders) {
  std::chrono::steady_clock::time_point t1 = std::chrono::steady_clock::now();

  ConstructAb(X, Y, orders, A_, b_);

  std::chrono::steady_clock::time_point t2 = std::chrono::steady_clock::now();
  std::chrono::microseconds time_used =
      std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1);
  printf("time used [%s] : %ld ms\n", __FUNCTION__, time_used.count());

  return SolveByLDLT(A_, b_);
}